import pandas as pd
import warnings
from TradeSignalCalculator import TradeSignalCalculator
from sqlalchemy import create_engine
from pathlib import Path
import numpy as np

warnings.filterwarnings('ignore')

class TradeProcessor:
    def __init__(self, db_connection_string=None, df=None, data_file_path=None):
        if db_connection_string:
            engine = create_engine(db_connection_string)
            query = "SELECT * FROM stock_basic_indicators"
            self.data = pd.read_sql(query, engine)
            engine.dispose()
        elif df is not None:
            self.data = df.copy()
        elif data_file_path:
            if not Path(data_file_path).exists():
                raise FileNotFoundError(f"Parquet文件未找到: {data_file_path}")
            self.data = pd.read_parquet(data_file_path)
            print(f"成功从 {data_file_path} 读取数据。")
        else:
            raise ValueError("必须提供db_connection_string、df或data_file_path参数")
        
        self._preprocess_data()
    
    def _preprocess_data(self):
        try:
            self.data['date'] = pd.to_datetime(self.data['date'], errors='coerce')
        except Exception as e:
            raise ValueError(f"日期转换错误: {e}")
        
        if self.data['date'].isna().sum() > 0:
            print(f"警告：存在无法解析的日期，已被转换为 NaT，数量：{self.data['date'].isna().sum()}")
            self.data = self.data.dropna(subset=['date'])
    
        self.data['date'] = self.data['date'].dt.normalize().dt.tz_localize(None)
        self.data['datekey'] = self.data['date'].dt.strftime('%Y%m%dd')
        
        self.data.set_index('date', inplace=True)
        
        # Clean price columns
        for col in ['close', 'low', 'high']:
            if col in self.data.columns:
                self.data = self.data[self.data[col] > 0]  # Remove non-positive prices
                self.data = self.data.dropna(subset=[col])  # Remove NaN
                self.data = self.data[~np.isinf(self.data[col])]  # Remove inf
                if self.data[col].isna().any() or np.isinf(self.data[col]).any():
                    print(f"Warning: After cleaning, {col} still contains invalid values.")
    
        self.data['action'] = 'na'
        self.data['sharpe_ratio'] = 0.0
        self.data['portfolio'] = 10000.0
        self.data['held'] = 0
    
        if 'code' not in self.data.columns:
            if isinstance(self.data.index, pd.MultiIndex):
                self.data.reset_index(level='code', inplace=True)
            else:
                raise ValueError("数据中缺少'code'列")     
            
    def apply_trade_logic(self, group):
            action_state = 'na'
            actions = []
            buy_signals = ['ENTRYBUY', 'BOTTOMBUY', 'BOTTOMUPBUY', 'POTENTIALBUY', 'STARK', 'DRAGONBUY']
            sell_signals = ['POWERDOWNSELL', 'CLEANSELL', 'STAGESELL']
    
            for index, row in group.iterrows():
                if action_state == 'na':
                    if any(row[sig] for sig in buy_signals) or (pd.notna(row['XYS1']) and pd.notna(row['XYS2']) and row['XYS1'] > row['XYS2']):
                        action_state = 'buy'
                        actions.append('buy')
                    else:
                        actions.append('na')
                elif action_state == 'hold':
                    if any(row[sig] for sig in sell_signals) or (pd.notna(row['XYS1']) and pd.notna(row['XYS2']) and row['XYS1'] < row['XYS2']):
                        action_state = 'sell'
                        actions.append('sell')
                    else:
                        actions.append('hold')
                elif action_state == 'buy':
                    action_state = 'hold'
                    actions.append('hold')
                elif action_state == 'sell':
                    action_state = 'na'
                    actions.append('na')
        
            group['action'] = actions
            return group
    
    def calculate_portfolio(self, group):
        portfolio_val = 10000.0  # 初始资金
        held_shares = 0
        entry_index = None
    
        portfolio_values = [0.0] * len(group)
        held_shares_list = [0] * len(group)
    
        holding_returns = []
    
        for i, (index, row) in enumerate(group.iterrows()):
            action = row['action']
            open_price = row['open']
            close_price = row['close']
    
            # 买入逻辑
            if action == 'buy' and held_shares == 0:
                max_shares = int(portfolio_val / open_price)
                shares_to_buy = (max_shares // 100) * 100 if max_shares >= 100 else 0
    
                if shares_to_buy > 0:
                    cost = shares_to_buy * open_price
                    portfolio_val -= cost
                    held_shares = shares_to_buy
                    entry_index = i
                    holding_returns = []
    
            # 记录持有期间的收益
            current_total = portfolio_val + held_shares * close_price
            if i > 0:
                prev_total = portfolio_values[i - 1]
                if prev_total > 0:
                    daily_return = (current_total - prev_total) / prev_total
                    holding_returns.append(daily_return)
            portfolio_values[i] = current_total
            held_shares_list[i] = held_shares
    
            # 卖出逻辑
            if action == 'sell' and held_shares > 0:
                sale_value = held_shares * open_price
                portfolio_val += sale_value
                held_shares = 0
                entry_index = None
                holding_returns = []
    
        # 更新 DataFrame
        group['portfolio'] = portfolio_values
        group['held'] = held_shares_list
        return group

    
    def process(self):
        df_temp = self.data.reset_index()

        processed_groups = []
        for code, group in df_temp.groupby('code'):
            group['date'] = pd.to_datetime(group['date'])
            group_for_calc = group.set_index('date').sort_index()

            calculator = TradeSignalCalculator(group_for_calc)
            calculated_group = calculator.calculate_all_signals()

            if 'action' not in calculated_group.columns:
                calculated_group['action'] = 'na'
            
            final_group_with_action = self.apply_trade_logic(calculated_group)
            
            # Now apply the Sharpe Ratio and portfolio logic
            final_group_with_sharpe = self.calculate_portfolio(final_group_with_action)
            
            final_group_with_sharpe = final_group_with_sharpe.reset_index()
            processed_groups.append(final_group_with_sharpe)
        
        self.data = pd.concat(processed_groups)
        self.data.set_index('date', inplace=True)
        self.data = self.data.sort_values(by=['code', 'date'])
        
        strategy_factors = ['date','code', 'close', 'open', 'low', 'high', 'volume','amount','turnover' ,'ddx','action' ,'POWERLINE', 'XYS1','XYS2','datekey', 'portfolio']
        available_factors = [f for f in strategy_factors if f in self.data.columns]
        self.data = self.data[available_factors]
        
        return self.data
    
    def save_result(self, output_path):
        self.data.to_parquet(output_path)

if __name__ == "__main__":
    current_script_dir = Path(__file__).parent
    input_file_path = current_script_dir / "G:/quant_data" / "basic_indicators.parquet"
    output_file_path = current_script_dir / "G:/quant_data" / "trade_indicators.parquet"
    class TradeSignalCalculator:
        def __init__(self, df_group):
            self.data = df_group.copy()
            for col in ['ENTRYBUY', 'BOTTOMBUY', 'BOTTOMUPBUY', 'POTENTIALBUY', 'STARK', 'DRAGONBUY',
                        'POWERDOWNSELL', 'CLEANSELL', 'STAGESELL', 'XYS1', 'XYS2']:
                if col not in self.data.columns:
                    self.data[col] = False if col in ['ENTRYBUY', 'BOTTOMBUY', 'BOTTOMUPBUY', 'POTENTIALBUY', 'STARK', 'DRAGONBUY',
                                                     'POWERDOWNSELL', 'CLEANSELL', 'STAGESELL'] else pd.NA
            # Example: fill some XYS values for testing
            self.data['XYS1'] = self.data['close'] * 1.01
            self.data['XYS2'] = self.data['close']
            # Example: Simulate some buy/sell signals
            self.data.loc[self.data.index.day == 5, 'ENTRYBUY'] = True
            self.data.loc[self.data.index.day == 15, 'CLEANSELL'] = True

        def calculate_all_signals(self):
            print("Running dummy TradeSignalCalculator.calculate_all_signals")
            return self.data.copy()

    processor = TradeProcessor(data_file_path=input_file_path)
    
    result = processor.process()  
    print("处理完成，前5行结果:")
    print(result.head())
    
    processor.save_result(output_file_path)
    print(f"结果已保存到: {output_file_path}")